"""
代码执行工具模块

提供安全的代码执行功能，支持多种编程语言
"""
import os
import subprocess
import sys
import tempfile
import time
from typing import Any, Dict, List, Optional

from .base_tool import BaseTool, ToolResult


class CodeExecutionTool(BaseTool):
    """
    代码执行工具
    
    提供安全的代码执行功能，支持Python、JavaScript、Shell等
    """
    
    def __init__(self, timeout: int = 30, max_memory: int = 100):
        """
        初始化代码执行工具
        
        Args:
            timeout: 执行超时时间（秒）
            max_memory: 最大内存使用量（MB）
        """
        super().__init__(
            name="code_execution_tool",
            description="执行代码，支持Python、JavaScript、Shell等多种语言",
        )
        
        self.timeout = timeout
        self.max_memory = max_memory
        self.supported_languages = {
            'python': self._execute_python,
            'py': self._execute_python,
            'javascript': self._execute_javascript,
            'js': self._execute_javascript,
            'shell': self._execute_shell,
            'bash': self._execute_shell,
            'sh': self._execute_shell,
            'powershell': self._execute_powershell,
            'ps1': self._execute_powershell
        }
        
        # 安全限制
        self.dangerous_modules = {
            'os', 'sys', 'subprocess', 'importlib', 'exec', 'eval',
            'compile', '__import__', 'open', 'file', 'input', 'raw_input'
        }
        
        self.allowed_modules = {
            'math', 'random', 'datetime', 'json', 're', 'collections',
            'itertools', 'functools', 'operator', 'string', 'time'
        }
    
    def execute(self, code: str, language: str, timeout: Optional[int] = None,
                input_data: Optional[str] = None, args: Optional[List[str]] = None,
                **kwargs) -> ToolResult:
        """
        执行代码
        
        Args:
            code: 要执行的代码
            language: 编程语言
            timeout: 超时时间
            input_data: 输入数据
            args: 命令行参数
            
        Returns:
            执行结果
        """
        try:
            # 验证语言支持
            if language.lower() not in self.supported_languages:
                return ToolResult(
                    success=False,
                    error_message=f"不支持的语言: {language}。支持的语言: {list(self.supported_languages.keys())}"
                )
            
            # 安全检查
            if language.lower() in ['python', 'py']:
                safety_check = self._check_python_safety(code)
                if not safety_check['safe']:
                    return ToolResult(
                        success=False,
                        error_message=f"代码安全检查失败: {safety_check['reason']}"
                    )
            
            # 执行代码
            executor = self.supported_languages[language.lower()]
            result = executor(code, timeout or self.timeout, input_data, args)
            
            return result
            
        except Exception as e:
            return ToolResult(
                success=False,
                error_message=f"代码执行失败: {str(e)}"
            )
    
    def _execute_python(self, code: str, timeout: int, input_data: Optional[str] = None,
                       args: Optional[List[str]] = None) -> ToolResult:
        """执行Python代码"""
        try:
            # 创建临时文件
            with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
                f.write(code)
                temp_file = f.name
            
            # 准备执行环境
            env = os.environ.copy()
            env['PYTHONPATH'] = os.pathsep.join(sys.path)
            
            # 执行代码
            cmd = [sys.executable, temp_file]
            if args:
                cmd.extend(args)
            
            process = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                stdin=subprocess.PIPE if input_data else None,
                text=True,
                env=env,
                cwd=tempfile.gettempdir()
            )
            
            # 发送输入数据
            if input_data:
                process.stdin.write(input_data)
                process.stdin.close()
            
            # 等待执行完成
            try:
                stdout, stderr = process.communicate(timeout=timeout)
                return_code = process.returncode
            except subprocess.TimeoutExpired:
                process.kill()
                stdout, stderr = process.communicate()
                return_code = -1
            
            # 清理临时文件
            try:
                os.unlink(temp_file)
            except Exception:
                pass
            
            # 处理结果
            if return_code == 0:
                return ToolResult(
                    success=True,
                    data={
                        "output": stdout,
                        "return_code": return_code
                    },
                    metadata={
                        "language": "python",
                        "execution_time": timeout,
                        "input_provided": input_data is not None
                    }
                )
            else:
                return ToolResult(
                    success=False,
                    error_message=f"执行失败 (返回码: {return_code}): {stderr}",
                    data={
                        "output": stdout,
                        "error": stderr,
                        "return_code": return_code
                    }
                )
                
        except Exception as e:
            return ToolResult(
                success=False,
                error_message=f"Python代码执行异常: {str(e)}"
            )
    
    def _execute_javascript(self, code: str, timeout: int, input_data: Optional[str] = None,
                           args: Optional[List[str]] = None) -> ToolResult:
        """执行JavaScript代码"""
        try:
            # 检查Node.js是否可用
            try:
                subprocess.run(['node', '--version'], capture_output=True, check=True)
            except (subprocess.CalledProcessError, FileNotFoundError):
                return ToolResult(
                    success=False,
                    error_message="Node.js未安装或不可用"
                )
            
            # 创建临时文件
            with tempfile.NamedTemporaryFile(mode='w', suffix='.js', delete=False) as f:
                f.write(code)
                temp_file = f.name
            
            # 执行代码
            cmd = ['node', temp_file]
            if args:
                cmd.extend(args)
            
            process = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                stdin=subprocess.PIPE if input_data else None,
                text=True,
                cwd=tempfile.gettempdir()
            )
            
            # 发送输入数据
            if input_data:
                process.stdin.write(input_data)
                process.stdin.close()
            
            # 等待执行完成
            try:
                stdout, stderr = process.communicate(timeout=timeout)
                return_code = process.returncode
            except subprocess.TimeoutExpired:
                process.kill()
                stdout, stderr = process.communicate()
                return_code = -1
            
            # 清理临时文件
            try:
                os.unlink(temp_file)
            except Exception:
                pass
            
            # 处理结果
            if return_code == 0:
                return ToolResult(
                    success=True,
                    data={
                        "output": stdout,
                        "return_code": return_code
                    },
                    metadata={
                        "language": "javascript",
                        "execution_time": timeout,
                        "input_provided": input_data is not None
                    }
                )
            else:
                return ToolResult(
                    success=False,
                    error_message=f"执行失败 (返回码: {return_code}): {stderr}",
                    data={
                        "output": stdout,
                        "error": stderr,
                        "return_code": return_code
                    }
                )
                
        except Exception as e:
            return ToolResult(
                success=False,
                error_message=f"JavaScript代码执行异常: {str(e)}"
            )
    
    def _execute_shell(self, code: str, timeout: int, input_data: Optional[str] = None,
                      args: Optional[List[str]] = None) -> ToolResult:
        """执行Shell脚本"""
        try:
            # 安全检查
            if not self._check_shell_safety(code):
                return ToolResult(
                    success=False,
                    error_message="Shell命令包含危险操作，执行被拒绝"
                )
            
            # 创建临时文件
            with tempfile.NamedTemporaryFile(mode='w', suffix='.sh', delete=False) as f:
                f.write("#!/bin/bash\n")
                f.write(code)
                temp_file = f.name
            
            # 设置执行权限
            os.chmod(temp_file, 0o755)
            
            # 执行脚本
            cmd = ['bash', temp_file]
            if args:
                cmd.extend(args)
            
            process = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                stdin=subprocess.PIPE if input_data else None,
                text=True,
                cwd=tempfile.gettempdir()
            )
            
            # 发送输入数据
            if input_data:
                process.stdin.write(input_data)
                process.stdin.close()
            
            # 等待执行完成
            try:
                stdout, stderr = process.communicate(timeout=timeout)
                return_code = process.returncode
            except subprocess.TimeoutExpired:
                process.kill()
                stdout, stderr = process.communicate()
                return_code = -1
            
            # 清理临时文件
            try:
                os.unlink(temp_file)
            except Exception:
                pass
            
            # 处理结果
            if return_code == 0:
                return ToolResult(
                    success=True,
                    data={
                        "output": stdout,
                        "return_code": return_code
                    },
                    metadata={
                        "language": "shell",
                        "execution_time": timeout,
                        "input_provided": input_data is not None
                    }
                )
            else:
                return ToolResult(
                    success=False,
                    error_message=f"执行失败 (返回码: {return_code}): {stderr}",
                    data={
                        "output": stdout,
                        "error": stderr,
                        "return_code": return_code
                    }
                )
                
        except Exception as e:
            return ToolResult(
                success=False,
                error_message=f"Shell脚本执行异常: {str(e)}"
            )
    
    def _execute_powershell(self, code: str, timeout: int, input_data: Optional[str] = None,
                           args: Optional[List[str]] = None) -> ToolResult:
        """执行PowerShell脚本"""
        try:
            # 检查PowerShell是否可用
            try:
                subprocess.run(['powershell', '-Command', 'Get-Host'], capture_output=True, check=True)
            except (subprocess.CalledProcessError, FileNotFoundError):
                return ToolResult(
                    success=False,
                    error_message="PowerShell未安装或不可用"
                )
            
            # 安全检查
            if not self._check_powershell_safety(code):
                return ToolResult(
                    success=False,
                    error_message="PowerShell命令包含危险操作，执行被拒绝"
                )
            
            # 执行PowerShell命令
            cmd = ['powershell', '-Command', code]
            if args:
                cmd.extend(args)
            
            process = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                stdin=subprocess.PIPE if input_data else None,
                text=True,
                cwd=tempfile.gettempdir()
            )
            
            # 发送输入数据
            if input_data:
                process.stdin.write(input_data)
                process.stdin.close()
            
            # 等待执行完成
            try:
                stdout, stderr = process.communicate(timeout=timeout)
                return_code = process.returncode
            except subprocess.TimeoutExpired:
                process.kill()
                stdout, stderr = process.communicate()
                return_code = -1
            
            # 处理结果
            if return_code == 0:
                return ToolResult(
                    success=True,
                    data={
                        "output": stdout,
                        "return_code": return_code
                    },
                    metadata={
                        "language": "powershell",
                        "execution_time": timeout,
                        "input_provided": input_data is not None
                    }
                )
            else:
                return ToolResult(
                    success=False,
                    error_message=f"执行失败 (返回码: {return_code}): {stderr}",
                    data={
                        "output": stdout,
                        "error": stderr,
                        "return_code": return_code
                    }
                )
                
        except Exception as e:
            return ToolResult(
                success=False,
                error_message=f"PowerShell脚本执行异常: {str(e)}"
            )
    
    def _check_python_safety(self, code: str) -> Dict[str, Any]:
        """检查Python代码安全性"""
        import ast
        
        try:
            tree = ast.parse(code)
        except SyntaxError as e:
            return {"safe": False, "reason": f"语法错误: {e}"}
        
        # 检查危险操作
        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                for alias in node.names:
                    if alias.name in self.dangerous_modules:
                        return {"safe": False, "reason": f"禁止导入模块: {alias.name}"}
            
            elif isinstance(node, ast.ImportFrom):
                if node.module in self.dangerous_modules:
                    return {"safe": False, "reason": f"禁止从模块导入: {node.module}"}
            
            elif isinstance(node, ast.Call):
                if isinstance(node.func, ast.Name):
                    if node.func.id in ['exec', 'eval', 'compile', '__import__']:
                        return {"safe": False, "reason": f"禁止调用函数: {node.func.id}"}
        
        return {"safe": True, "reason": "代码安全检查通过"}
    
    def _check_shell_safety(self, code: str) -> bool:
        """检查Shell命令安全性"""
        dangerous_commands = [
            'rm -rf', 'rm -r', 'rm -f', 'rm /', 'rm /*',
            'format', 'fdisk', 'mkfs', 'dd if=',
            'wget', 'curl', 'nc', 'netcat', 'telnet',
            'ssh', 'scp', 'rsync', 'ftp',
            'sudo', 'su', 'chmod 777', 'chown',
            'killall', 'pkill', 'kill -9',
            'shutdown', 'reboot', 'halt', 'poweroff'
        ]
        
        code_lower = code.lower()
        for cmd in dangerous_commands:
            if cmd in code_lower:
                return False
        
        return True
    
    def _check_powershell_safety(self, code: str) -> bool:
        """检查PowerShell命令安全性"""
        dangerous_commands = [
            'Remove-Item', 'rm', 'del', 'format',
            'Invoke-WebRequest', 'iwr', 'wget',
            'Invoke-Expression', 'iex', 'eval',
            'Start-Process', 'Start-Service', 'Stop-Service',
            'Set-ExecutionPolicy', 'Get-Content', 'Set-Content',
            'New-Item', 'Copy-Item', 'Move-Item'
        ]
        
        code_lower = code.lower()
        for cmd in dangerous_commands:
            if cmd.lower() in code_lower:
                return False
        
        return True
    
    def get_supported_languages(self) -> List[str]:
        """获取支持的语言列表"""
        return list(self.supported_languages.keys())
    
    def get_safety_info(self) -> Dict[str, Any]:
        """获取安全信息"""
        return {
            "dangerous_modules": list(self.dangerous_modules),
            "allowed_modules": list(self.allowed_modules),
            "timeout": self.timeout,
            "max_memory": self.max_memory,
            "supported_languages": self.get_supported_languages()
        }


class CodeExecutor:
    """代码执行器管理器"""
    
    def __init__(self):
        """初始化代码执行器"""
        self.tools = {}
        self.execution_history = []
    
    def get_tool(self, name: str = "default") -> CodeExecutionTool:
        """
        获取代码执行工具
        
        Args:
            name: 工具名称
            
        Returns:
            代码执行工具实例
        """
        if name not in self.tools:
            self.tools[name] = CodeExecutionTool()
        
        return self.tools[name]
    
    def execute_code(self, code: str, language: str, **kwargs) -> ToolResult:
        """
        执行代码的便捷方法
        
        Args:
            code: 代码
            language: 语言
            **kwargs: 其他参数
            
        Returns:
            执行结果
        """
        tool = self.get_tool()
        result = tool.execute(code, language, **kwargs)
        
        # 记录执行历史
        self.execution_history.append({
            "timestamp": time.time(),
            "language": language,
            "code": code[:100] + "..." if len(code) > 100 else code,
            "success": result.success
        })
        
        # 限制历史记录数量
        if len(self.execution_history) > 100:
            self.execution_history = self.execution_history[-100:]
        
        return result
    
    def get_execution_history(self, limit: int = 10) -> List[Dict]:
        """获取执行历史"""
        return self.execution_history[-limit:]


# 全局代码执行器
_code_executor = CodeExecutor()


def get_code_execution_tool(name: str = "default") -> CodeExecutionTool:
    """获取代码执行工具的便捷函数"""
    return _code_executor.get_tool(name)


def execute_code(code: str, language: str, **kwargs) -> ToolResult:
    """执行代码的便捷函数"""
    return _code_executor.execute_code(code, language, **kwargs)
